from flask import request
from flask_restful import Resource, fields, marshal_with
from models import  Teacher, Student
from db.connection import db

teacher_fields = {
    'id': fields.Integer,
    'name': fields.String,
    'subject': fields.String,
}

student_fields = {
    'id': fields.Integer,
    'name': fields.String,
    'age': fields.Integer,
    'grade': fields.String,
    'teacher_id': fields.Integer,
}

class TeacherList(Resource):
    @marshal_with(teacher_fields)
    def get(self):
        teachers = Teacher.query.all()
        return teachers

    def post(self):
        data = request.get_json()
        new_teacher = Teacher(name=data['name'], subject=data['subject'])
        db.session.add(new_teacher)
        db.session.commit()
        return {'message': 'Teacher created successfully'}, 201

class TeacherDetail(Resource):
    @marshal_with(teacher_fields)
    def get(self, teacher_id):
        teacher = Teacher.query.get_or_404(teacher_id)
        return teacher

    def put(self, teacher_id):
        data = request.get_json()
        teacher = Teacher.query.get_or_404(teacher_id)
        teacher.name = data['name']
        teacher.subject = data['subject']
        db.session.commit()
        return {'message': 'Teacher updated successfully'}, 200

    def delete(self, teacher_id):
        teacher = Teacher.query.get_or_404(teacher_id)
        db.session.delete(teacher)
        db.session.commit()
        return {'message': 'Teacher deleted successfully'}, 200

class StudentList(Resource):
    @marshal_with(student_fields)
    def get(self):
        students = Student.query.all()
        return students

    def post(self):
        data = request.get_json()
        new_student = Student(name=data['name'], age=data['age'], grade=data['grade'], teacher_id=data['teacher_id'])
        db.session.add(new_student)
        db.session.commit()
        return {'message': 'Student created successfully'}, 201

class StudentDetail(Resource):
    @marshal_with(student_fields)
    def get(self, student_id):
        student = Student.query.get_or_404(student_id)
        return student

    def put(self, student_id):
        data = request.get_json()
        student = Student.query.get_or_404(student_id)
        student.name = data['name']
        student.age = data['age']
        student.grade = data['grade']
        student.teacher_id = data['teacher_id']
        db.session.commit()
        return {'message': 'Student updated successfully'}, 200

    def delete(self, student_id):
        student = Student.query.get_or_404(student_id)
        db.session.delete(student)
        db.session.commit()
        return {'message': 'Student deleted successfully'}, 200
